"""
中心词 + 上下文
"""
import re

import torch

from dataset.vocab import Vocab, read_txt, tokenizer
from torch.utils.data import Dataset, DataLoader

tokens = tokenizer(read_txt("./data/timemachine.txt"))
vocab = Vocab(tokens, 0)
token_ids = [vocab.to_idx(line) for line in tokens]


def _generate_center_and_context(token_ids, window_size=2):
    centers = []  # 记录中心词
    contexts = []  # 记录对应中心词的上下文
    # 获取每一个词
    for line in token_ids:
        for i, idx in enumerate(line):
            start = max(0, i - window_size)
            end = min(len(line), i + window_size + 1)
            context = line[start:i] + line[i + 1:end]
            centers += [idx] * len(context)
            contexts += context
    return torch.tensor(centers, dtype=torch.long), torch.tensor(contexts, dtype=torch.long)


# centers, contexts = _generate_center_and_context(token_ids)

class SkipGramDataset(Dataset):
    def __init__(self):
        super().__init__()
        self.centers, self.contexts = _generate_center_and_context(token_ids)

    def __len__(self):
        return len(self.centers)

    def __getitem__(self, index):
        return self.centers[index], self.contexts[index]


def _generate_loader(batch_size=100):
    dataset = SkipGramDataset()
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return dataloader, vocab